import os, sys
import numpy as np
from matplotlib.mlab import find, prctile
from matplotlib.figure import Figure
from matplotlib.patches import Rectangle
from matplotlib.ticker import MultipleLocator
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from cwd.lib.data import MonthlyGrid
from cwd.lib.utils import length, TextOutput

class MonthlyPercentile(object):
    def __init__(self, lat, lon, frztype, units='m', data=None):
        self.lat = lat
        self.lon = lon
        self.frztype = frztype
        self.units = units
        if data is None:
            grid = MonthlyGrid(lat,lon,frztype,units)
            self.data = grid.data
            self.last_day = grid.last_day
            
    def write_to(self, output, format="png"):
        """
        Writes image to file-like 'output'
        """
        if format == 'text':
            # Pass all parameters to graphing function
            self.graph_ax(output, format)
            output.get_text()
        else:
            # Create figure (the plot container)
            fig = Figure(figsize=(12,7), facecolor='w')
            # Pass all parameters to graphing function (alters 'fig')
            self.graph_ax(fig, format)
            # Create a canvas for png write
            canvas = FigureCanvas(fig)
            canvas.print_png(output)

    def graph_ax(self, outobj, format):
        """
        pass in the latitude, longitude requested
        also pass in the fixed temperature surface requested
        0:0C, 1:10C, 2:20C, 3:30C
        This function reads in all monthly data and performs
        inverse distance weighting to interpolate to given lat/lon
        Plots out time series of last year of data, climatology, and the
        threshold values
        """

        lat = self.lat
        lon = self.lon
        frztype = self.frztype
        DATA = self.data

        frztype_map = { 0: 0, 1: 10, 2: 20, 3: 30, }
        perc_colors = ( [.85, .99, .99], [.7, .99, .99], [.5, .99, .99] )
        plevels = [5., 10., 25., 50., 75., 90., 95.]

        # this checks to see how many months have actual data, generalized to account
        # for having missing data in the array (e.g., May 2009, having missing data for Jun-Dec)
        lastmo = length(find(np.logical_not(np.isnan(DATA.ravel()))))
        lastyr = int(np.ceil(lastmo/12.))
        lastmo = np.mod(lastmo,12)
        # lastmo is 1-12
        if lastmo == 0: lastmo = 12

        # gauge percentile values w/respect to the 62 yr record
        # there are 7 levels of interest to examine
        # loop over months
        perc = np.zeros((12,7))
        for i in range(12):
            for j in range(7):
                perc[i][j] = prctile(DATA[i,:], plevels[j]) # perc(i,j)=prctile(DATA(i,:),plevels(j));

        # this assumes currentdayofyear is not 366
        # the goal here is to create a plot that shows the last 12 months w/respect to
        # climatology, so we need to figure out what the last twelve months are
        # since this requires breaking the dataset in groups (unless we have Jan-Dec),
        # four variables are declared, d1 starts with the month 12 months ago, goes until Dec (d2)
        # then d3 starts in Jan and ends in current (or last) month d4
        d3 = 1
        d4 = lastmo # d4 is the last month
        d2 = 12
        d1 = d4+1 # d3 represents the first month, or 11 months ago
        if d3 > 12: d3=1

        # plotting
        dataorder = np.zeros(24)
        dataorder[:12] = np.arange(0,12) # dataorder=1:12
        dataorder[12:24] = np.arange(11,-1,-1) # dataorder(13:24)=12:-1:1;
        # Setup output object
        if format != 'text':
            ax = outobj.add_axes([0.1, 0.05, 0.6, .85])

        # matlabs patch function looks similar to the fill_between pylab plot
        # the matlab patch function requires essentially for you to draw a rectangle, that is
        # we have an index, dataorder, going from 1:12 and from 12:-1:1, so this is where the d1 d2 d3 d4
        # comes into play in splitting up the dataset
        if lastmo == 12: # if lastmo==12, no need for d{1..4}
            # Take first and last pairs, moving inward
            # plevels = [5., 10., 25., 50., 75., 90., 95.]
            # [0, 1, 2, 3, 4, 5, 6]
            # [(0,6), (1,5), (2,4)]
            for i,(y1,y2) in enumerate([(0,6), (1,5), (2,4)]):
                ext = np.zeros(dataorder.shape)
                plen = length(perc[:,y1])
                ext[:plen] = perc[:,y1] # ext=perc(:,y1);
                ext[plen:] = perc[::-1,y2] # ext(13:24)=perc(12:-1:1,y2);
                if format == 'text':
                    outobj.add(perc[:,y1], header="%dth" % plevels[y1])
                    outobj.add(perc[:,y2], header="%dth" % plevels[y2])
                else:
                    ax.fill_between(dataorder, ext, color=perc_colors[i], edgecolor='w')
                    # a=patch(dataorder,ext,[.85 .99 .99]);set(a,'edgecolor','w')

            # next plot median value
            datam = perc[:,3] # datam=perc(:,4);
            if format == 'text':
                outobj.add(datam, header="median")
            else:
                meanline = ax.plot(datam, 'k', linewidth=2) # a=plot(datam,'k');set(a,'linewidth',2);

            # finally plot last 12 months
            datam = DATA[:,lastyr-1] # datam=DATA(:,lastyr);
        else: # lastmo < 12
            # Take first and last pairs, moving inward
            # [0, 1, 2, 3, 4, 5, 6]
            # [(0,6), (1,5), (2,4)]
            for i,(y1,y2) in enumerate([(0,6), (1,5), (2,4)]):
                ext = np.zeros(dataorder.shape)
                plen = length(perc[d1-1:d2,y1])
                ext[:plen] = perc[d1-1:d2,y1] # ext=perc(d1:d2,y1)
                ext[plen:12] = perc[d3-1:d4,y1] # ext(length(ext)+1:12)=perc(d3:d4,y1);   
                ext[12:13+d4-d3] = perc[np.arange(d4-1,d3-1-1,-1),y2] # ext(13:13+d4-d3)=perc(d4:-1:d3,y2)
                ext[13+d4-d3:24] = perc[np.arange(d2-1,d1-1-1,-1),y2] # ext(13+d4-d3+1:24)=perc(d2:-1:d1,y2);
                if format == 'text':
                    outobj.add(ext[:12], header="%dth"%plevels[y1])
                    outobj.add(ext[::-1][:12], header="%dth"%plevels[y2])
                else:
                    ax.fill_between(dataorder, ext, color=perc_colors[i], edgecolor='w')
                    # a=patch(dataorder,ext,[.5 .99 .99]);set(a,'edgecolor','w');

            # next plot median value
            datam = np.zeros(12)
            plen = length(perc[d1-1:d2,3])
            datam[:plen] = perc[d1-1:d2,3] # datam=perc(d1:d2,4)
            datam[plen:] = perc[d3-1:d4,3] # datam(length(datam)+1:12)=perc(d3:d4,4);
            if format == 'text':
                outobj.add(datam, header="median")
            else:
                meanline = ax.plot(datam, 'k', linewidth=2) # a=plot(datam,'k');set(a,'linewidth',2);

            # finally plot last year
            datam = np.zeros(12)
            plen = length(DATA[d1-1:d2,lastyr-1-1])
            datam[:plen] = DATA[d1-1:d2,lastyr-1-1] # datam=DATA(d1:d2,lastyr-1)
            datam[plen:] = DATA[d3-1:d4,lastyr-1] # datam(length(datam)+1:12)=DATA(d3:d4,lastyr);

        if format == 'text':
            outobj.add(datam, header="Last 12")
        else:
            MONTH_ABBR = ['J','F','M','A','M','J','J','A','S','O','N','D']
            ax.set_xticks(np.arange(12))
            ax.set_xticklabels(MONTH_ABBR[lastmo:] + MONTH_ABBR[:lastmo-1] + [self.last_day.strftime("%b %d"),])

            # Plot last 12 Months
            last12line = ax.plot(datam, 'r', linewidth=2) # a=plot(datam,'r');set(a,'linewidth',2);
            ax.plot(datam, 'rx', linewidth=2, markersize=18) # a=plot(datam,'rx');set(a,'linewidth',2,'markersize',18);

            # Set legend
            labels = ( r"$5-95^{th}$", r"$10-90^{th}$", r"$25-75^{th}$", 'Median', 'Last 12 Months')
            outobj.legend( [ Rectangle((0,0),1,1, facecolor=c, edgecolor='w') for c in perc_colors ]+[meanline,last12line],
                        labels, loc=(.71,.665),shadow=True)

            # Set title
            # title(['Elevation of 0\circC Surface for ',lats,'\circN, ',lons,'\circW'],'fontsize',28)
            ax.set_title(u"Elevation of %d\u00b0C Surface for %4.2f\u00b0N, %4.2f\u00b0W" % (frztype_map[frztype],lat,lon),
                        fontsize='xx-large', family='serif')
            ax.set_ylabel('Elevation (%s)' % self.units) # ylabel('Elevation (m)','fontsize',18)  
            ax.autoscale_view(tight=True) # axis tight;
            ax.yaxis.set_major_locator(MultipleLocator(500))
            ax.yaxis.set_minor_locator(MultipleLocator(100))
            outobj.text(0.71, 0.60, "Western Regional\nClimate Center",
                        {'size':"x-small", 'family':"monospace"})

if __name__ == '__main__':
    #output = open('test_monthlypercentile.png','wb')
    #graph = MonthlyPercentile(39,119,0)
    #graph.write_to(output)
    #output.close()
    text_output = TextOutput(delim=",")
    graph = MonthlyPercentile(40,117.5,0,units='m')
    graph.write_to(text_output,format="text")
